1102D - Balanced Ternary String - CodeForces Solution

//
//  main.cpp
//  test
//
//  Created by Ziheng Bao on 2023/4/15.
//

#include <bits/stdc++.h>

using namespace std;

class A{
public:
    const static void func(){
        cout << "func" << endl;
    }
};

void _531(){
    int n;
    cin >> n;
    vector<int> a(n);
    for(int i=0;i<n;i++){
        cin >> a[i];
    }
    vector<long long> pre(n+1, 0);
    long long now = 0;
    for(int i=0;i<n;i++){
        now += a[i];
        pre[i+1] = now;
    }
    int h = 0;
    deque<int> q;
    for(int i=0;i<=n;i++){
        while(!q.empty() && pre[i] > pre[q.back()]){
            q.pop_back();
        }
        q.push_back(i);
    }
    int x = 0, y = 0, z = 0;
    for(int i=0;i<=n;i++){
        if(pre[i] > pre[h]){
            h = i;
        }
        if(q.front() < i){
            q.pop_front();
        }
        if(pre[i]-pre[h] - pre[q.front()] < pre[x]-pre[y]-pre[z]){
            x = i;
            y = h;
            z = q.front();
        }
    }
    cout << y << " " <<  x << " " << z << endl;
}

void _529(){
    int n;
    cin >> n;
    string s;
    cin >> s;
    int c0 = 0, c1 = 0, c2 = 0;
    for(auto c:s){
        if(c == '0')
            c0++;
        else if(c == '1')
            c1++;
        else
            c2++;
    }
    for(auto& c:s){
        if(c == '2' && c2 > n/3){
            if(c0 < n/3){
                c = '0';
                c2--;
                c0++;
            }
            else if(c1 < n/3){
                c = '1';
                c2--;
                c1++;
            }
        }
        else if(c == '1' && c1 > n/3){
            if(c0 < n/3){
                c = '0';
                c1--;
                c0++;
            }
        }
    }
    for(int i=n-1;i>=0;i--){
        auto& c = s[i];
        if(c == '0' && c0 > n/3){
            if(c2 < n/3){
                c = '2';
                c0--;
                c2++;
            }
            else if(c1 < n/3){
                c = '1';
                c0--;
                c1++;
            }
        }
        else if(c == '1' && c1 > n/3){
            if(c2 < n/3){
                c = '2';
                c1--;
                c2++;
            }
        }
    }
    cout << s << endl;
}

int main(int argc, const char * argv[]) {
    _529();
}